package org.example.myleet.p648;

import java.util.List;

public class Solution {
    public String replaceWords(List<String> dictionary, String sentence) {
        MyTrieTree trieTree = new MyTrieTree();
        for (String word : dictionary) {
            trieTree.add(word);
        }
        String[] words = sentence.split(" ");
        StringBuilder result = new StringBuilder();
        for (String word : words) {
            String prefix = trieTree.findPrefix(word);
            if (prefix.length() > 0) {
                result.append(prefix).append(' ');
            } else {
                result.append(word).append(' ');
            }
        }
        result.deleteCharAt(result.length() - 1);
        return result.toString();
    }

    static class MyTrieTree implements TrieTree {
        TrieNode root;

        public MyTrieTree() {
            root = new TrieNode('0');
        }

        @Override
        public void add(String word) {
            TrieNode cur = root;
            for (int i = 0; i < word.length(); ++i) {
                char c = word.charAt(i);
                int idx = c - 'a';
                if (null == cur.children[idx]) {
                    cur.children[idx] = new TrieNode(c);
                }
                cur = cur.children[idx];
            }
            cur.end = true;
        }

        @Override
        public String findPrefix(String word) {
            StringBuilder sb = new StringBuilder();
            boolean found = true;
            TrieNode cur = root;
            for (int i = 0; i < word.length(); ++i) {
                char c = word.charAt(i);
                int idx = c - 'a';
                if (null == cur.children[idx]) {
                    found = false;
                    break;
                }
                sb.append(c);
                cur = cur.children[idx];
                if (cur.end) {
                    break;
                }
            }
            if (found) {
                return sb.toString();
            }
            return "";
        }
    }

    static class TrieNode {
        TrieNode[] children;
        boolean end;
        char c;

        public TrieNode(char c) {
            this.c = c;
            end = false;
            children = new TrieNode[26];
        }
    }

    interface TrieTree {

        /**
         * 向字典树添加一个单词
         */
        void add(String word);

        /**
         * 在字典树中寻找是否存在word的前缀，如果存在则返回前缀，不存在则返回空字符串
         */
        String findPrefix(String word);
    }
}
